import jax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array, Float, Int, Bool, PyTree
from .modules.utils import forward_pass, get_dts # , pad_to_len
from .modules.ode import integrate

# Carry: c, c_bar, delta, gate_o
# Output: h_t
# Update step: decay(c, c_bar, delta, dt)->ct ouput(ct, gate_o)->ht jump(x, h, c, c_bar)->c, c_bar, delta, gate_o

class CTLSTMCell(eqx.Module):
    stack: eqx.Module

    def __init__(self, hdim: int, input_dim: int, key: Array):
        self.stack = eqx.nn.Linear(hdim+input_dim, 7*hdim, key=key)

    def __call__(self, x: Float[Array, "input_dim"], h: Float[Array, "hdim"], c: Float[Array, "hdim"], c_bar: Float[Array, "hdim"]):
        hdim = h.shape[-1]
        cat = jnp.concatenate((x, h)) # (input_dim+hdim, )
        param = self.stack(cat)
        boundary = 6*hdim
        param1 = param[:boundary]
        param2 = param[boundary:]
        gate_i, gate_f, z, gate_o, gate_i_bar, gate_f_bar = jnp.split(jax.nn.sigmoid(param1), 6, -1)
        z = 2*z - 1
        delta = jax.nn.softplus(param2)
        c = gate_f * c + gate_i * z
        c_bar = gate_f_bar * c_bar + gate_i_bar * z
        return c, c_bar, delta, gate_o

def decay(c: Float[Array, "hdim"], c_bar: Float[Array, "hdim"], delta: Float[Array, "hdim"], dtime: Float[Array, ""]):
    c_t = c_bar + (c - c_bar) * jnp.exp(-delta * dtime)
    return c_t # (hdim, )

def output(ct: Float[Array, "dim"], gate_o: Float[Array, "dim"]):
    h_t = gate_o * (2*jax.nn.sigmoid(2*ct)-1)
    return h_t

class CTLSTM(eqx.Module):
    cell: CTLSTMCell
    hdim: int

    def __init__(self, hdim: int, embed_dim: int, key: Array):
        self.cell = CTLSTMCell(hdim, embed_dim, key)
        self.hdim = hdim

    def __call__(self, dts: Float[Array, "T"], embeds: Float[Array, "T embed_dim"]):
        carry_init = tuple(jnp.split(jnp.zeros(4*self.hdim, dtype=float), 4))
        xs = (dts, embeds)
        def f(carry, x):
            dt, embed = x
            carry = self.step(carry, embed, dt)
            return carry, carry
        carry, ys = jax.lax.scan(f, carry_init, xs)
        return carry, ys
    

    def step(self, carry: PyTree, embed: Float[Array, "embed_dim"], dt: Float[Array, ""]):
        c, c_bar, delta, gate_o = carry
        c_t = decay(c, c_bar, delta, dt)
        h_t = output(c_t, gate_o)
        c, c_bar, delta, gate_o = self.cell(embed, h_t, c, c_bar)
        return c, c_bar, delta, gate_o

class NHP(eqx.Module):
    embed: eqx.nn.Embedding
    intensity_layer: list
    enc: CTLSTM

    def __init__(self, hdim: int, num_types: int, embed_dim: int, key: Array):
        key1, key2, key3 = jax.random.split(key, 3)
        self.enc = CTLSTM(hdim, embed_dim, key1)
        self.embed = eqx.nn.Embedding(num_types, embed_dim, key=key2)
        self.intensity_layer = [
            eqx.nn.Linear(hdim, num_types, key=key3),
            jax.nn.softplus
        ]

    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        _, states = self.enc(dts, embeds) # tuple of 4 (T, hdim)s
        ll, dt_ll, mark_ll = jax.vmap(self._get_ll)(
            jax.tree.map(lambda vec: vec[:-1], states),
            dts[1:],
            marks[1:],
            mask[1:]
        )
        ll = ll.sum()
        dt_ll = dt_ll.sum()
        mark_ll = mark_ll.sum()
        return ll, dt_ll, mark_ll

    def _get_ll(self, state: PyTree, dt: Float[Array, ""], mark: Int[Array, ""], mask: Bool[Array, ""]):
        sample_ratios = jnp.linspace(0., 1., 100) # (100, )
        sampled_dts = sample_ratios * dt # (100, )
        intensities = jax.vmap(self._get_intensities, (None, 0))(state, sampled_dts) # (100, num_types)
        non_intensity = intensities[:-1].sum(-1) # (99, )
        Lambda = non_intensity.mean() * dt # ()
        event_intensities = intensities[-1] # (num_types)
        event_intensity = event_intensities[mark] # ()
        ll = jnp.log(event_intensity+1e-6) - Lambda
        mark_dist = event_intensity / (event_intensities.sum()+1e-6)
        mark_ll = jnp.log(mark_dist+1e-6)
        dt_ll = ll - mark_ll
        ll = jnp.where(mask, ll, 0.)
        dt_ll = jnp.where(mask, dt_ll, 0.)
        mark_ll = jnp.where(mask, mark_ll, 0.)
        return ll, dt_ll, mark_ll

    def _get_intensities(self, state: PyTree, dt: Float[Array, ""]):
        c, c_bar, delta, gate_o = state
        ct = decay(c, c_bar, delta, dt)
        ht = output(ct, gate_o) # (hdim, )
        intensities = forward_pass(self.intensity_layer, ht) # (num_types, )
        return intensities

    def intensity_at(self, state: PyTree, dt: Float[Array, ""]):
        intensities = self._get_intensities(state, dt)
        intensity = intensities.sum()
        return intensity

    def _predict(self, state: PyTree, dt_max: Float[Array, ""]): # c, c_bar, delta, gate_o
        Lambda0 = jnp.asarray(0., dtype=jnp.float32)
        int_tf0 = jnp.asarray(0., dtype=jnp.float32)
        H0 = (Lambda0, int_tf0)
        def func(t, H, args=None):
            Lambda, int_tf = H
            intensity = self.intensity_at(state, t) # ()
            f_dt = intensity * jnp.exp(-Lambda)
            tf = t * f_dt
            return intensity, tf
        _, Efdt = integrate(func, 0., dt_max, H0, None)
        intensities = self._get_intensities(state, Efdt) # (num_types, )
        mark_predict = intensities.argmax(-1)
        return Efdt, mark_predict
    
    @eqx.filter_jit
    def rolling_predict(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], dt_max: float, key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        _, states = self.enc(dts, embeds)
        dt_predict, mark_predict = jax.vmap(self._predict, (0, None))\
            (jax.tree.map(lambda vec: vec[:-1], states), dt_max) # (T-1, )
        return (dt_predict, mark_predict), (dts[1:], marks[1:]), mask[1:]

    @eqx.filter_jit
    def encode(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key=None):
        dts = get_dts(ts)
        dts = jnp.where(mask, dts, 0.)
        embeds = jax.vmap(self.embed)(marks)
        _, states = self.enc(dts, embeds) # tuple of 4 (T, hdim)s
        return states
    
    @eqx.filter_jit
    def intensities_at(self, state: PyTree, dts: Float[Array, "num_samples"], key=None):
        intensities = jax.vmap(self._get_intensities, (None, 0))(state, dts) # (num_samples, num_types)
        return intensities